from contextlib import asynccontextmanager
from typing import AsyncGenerator

from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from sqlalchemy.ext.declarative import declarative_base
from sqlmodel import SQLModel

from src.config.database import mysqlConfig

# 异步数据库连接配置
if mysqlConfig.ASYNC_DATABASE_URL:
    ASYNC_SQLALCHEMY_DATABASE_URL = mysqlConfig.ASYNC_DATABASE_URL
else:
    ASYNC_SQLALCHEMY_DATABASE_URL = f"mysql+aiomysql://{mysqlConfig.DB_USER}:{mysqlConfig.DB_PASSWORD}@{mysqlConfig.DB_HOST}:{mysqlConfig.DB_PORT}/{mysqlConfig.DB_NAME}"

# 异步引擎和会话
async_engine = create_async_engine(ASYNC_SQLALCHEMY_DATABASE_URL, **mysqlConfig.ASYNC_SQLALCHEMY_ENGINE_OPTIONS)
AsyncSessionLocal = async_sessionmaker(
    bind=async_engine,
    class_=AsyncSession,
    autocommit=False,
    autoflush=False,
    expire_on_commit=False  # 避免在commit后访问属性时出错
)

Base = declarative_base()


async def create_db_and_tables_async():
    """创建数据库表（异步）"""
    async with async_engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)


@asynccontextmanager
async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
    """异步数据库会话"""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()
